import cv2
import face_recognition
import numpy as np
import os

'''
@文件描述：读取摄像头rtsp视频流,进行人脸比
@创建人：wanghua
'''

class faceRecognition:
    def __init__(self, face_path='/home/hdkx/software/faceImage', rtsp_url='rtsp://admin:admin123@192.168.26.222:554'
                                                                           '/cam/realmonitor?channel=1&subtype=0'):
        self.face_path = face_path
        self.rtsp_url = rtsp_url
        self.person_encodings = []
        self.person_names = []
        self.getKnownFaceEncodings()

    def start(self):
        cap = cv2.VideoCapture(self.rtsp_url)
        while True:
            rec, frame = cap.read()
            # 将视频帧大小调整为1/4，以加快人脸识别处理
            small_frame = cv2.resize(frame, (0, 0), fx=0.25, fy=0.25)
            # 将图像从BGR颜色（OpenCV使用）转换为RGB颜色（face\u recognition使用）
            rgb_small_frame = small_frame[:, :, ::-1]
            # 查找当前视频帧中的所有人脸
            face_locations = face_recognition.face_locations(rgb_small_frame)
            # 获得当前视频帧中的人脸编码
            face_encodings = face_recognition.face_encodings(rgb_small_frame, face_locations)
            face_names = []
            #  匹配人脸
            for face_encoding in face_encodings:
                # 查看该人脸是否与已知人脸匹配
                matches = face_recognition.compare_faces(self.person_encodings, face_encoding, tolerance=0.5)
                # 或者，使用与新人脸距离最小的已知人脸
                face_distances = face_recognition.face_distance(self.person_encodings, face_encoding)
                best_match_index = np.argmin(face_distances)
                if matches[best_match_index]:
                    name = self.person_names[best_match_index]
                face_names.append(name)
            print(face_names)

    def getKnownFaceEncodings(self):
        # 读取人脸库文件夹
        face_ist = os.listdir(self.face_path)
        for i, each in enumerate(face_ist):
            image = face_recognition.load_image_file(f"{self.face_path}/{each}")
            face_encoding = face_recognition.face_encodings(image)[0]
            self.person_encodings.append(face_encoding)
            self.person_names.append(each)